# 创建密钥矩阵
def create_matrix(key):
   key = ''.join(sorted(set(key), key=key.index)).replace('J', 'I')
   alphabet = "ABCDEFGHIKLMNOPQRSTUVWXYZ"
   matrix = [c for c in key if c in alphabet]
   matrix += [c for c in alphabet if c not in matrix]
   return [matrix[i:i+5] for i in range(0, 25, 5)]

def out_matrix(matrix):
    print('密钥矩阵')
    for i in matrix:
        print(i)
# 获取字符在矩阵中的位置
def get_position(char, matrix):
   for i, row in enumerate(matrix):
       if char in row:
           return i, row.index(char)
   return None
# 加密字符对
def encrypt_pair(a, b, matrix):
   r1, c1 = get_position(a, matrix)
   r2, c2 = get_position(b, matrix)
   if r1 == r2:
       return matrix[r1][(c1 + 1) % 5] + matrix[r2][(c2 + 1) % 5]
   elif c1 == c2:
       return matrix[(r1 + 1) % 5][c1] + matrix[(r2 + 1) % 5][c2]
   else:
       return matrix[r1][c2] + matrix[r2][c1]
# 解密字符对
def decrypt_pair(a, b, matrix):
   r1, c1 = get_position(a, matrix)
   r2, c2 = get_position(b, matrix)
   if r1 == r2:
       return matrix[r1][(c1 - 1) % 5] + matrix[r2][(c2 - 1) % 5]
   elif c1 == c2:
       return matrix[(r1 - 1) % 5][c1] + matrix[(r2 - 1) % 5][c2]
   else:
       return matrix[r1][c2] + matrix[r2][c1]
# 加密函数
def playfair_encrypt(plaintext, key):
   plaintext = plaintext.upper().replace('J', 'I').replace(' ', '')
   if len(plaintext) % 2 != 0:
       plaintext += 'X'
   pairs = [plaintext[i:i+2] for i in range(0, len(plaintext), 2)]
   matrix = create_matrix(key)
   out_matrix(matrix)
   ciphertext = ''.join([encrypt_pair(a, b, matrix) for a, b in pairs])
   return ciphertext
# 解密函数
def playfair_decrypt(ciphertext, key):
   pairs = [ciphertext[i:i+2] for i in range(0, len(ciphertext), 2)]
   matrix = create_matrix(key)
   plaintext = ''.join([decrypt_pair(a, b, matrix) for a, b in pairs])
   return plaintext
# 示例
key = "ZHANGYIXINZ"
plaintext = "I like applied cryptography"
ciphertext = playfair_encrypt(plaintext, key)
print(f"加密后的密文：{ciphertext}")
#decrypted_text = playfair_decrypt(ciphertext, key)
#print(f"解密后的明文：{decrypted_text}")
